# -*- coding = utf-8 -*-
# @Time : 2022/1/15 21:19
# @Author : Chunyan Wei
# @File : train.py
# @Software:PyCharm

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pylab as plt
import numpy as np

def main():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
    #50000张训练图片
    train_set = torchvision.datasets.CIFAR10(root='./CIFAR10data',train=True,
                                             download=False,transform=transform)
    train_loader = torch.utils.data.DataLoader(train_set,batch_size=36,
                                               shuffle=True,num_workers=0)
    #10000张验证图片
    val_set = torchvision.datasets.CIFAR10(root='./CIFAR10data',train=False,
                                           download=False,transform=transform)
    val_loader = torch.utils.data.DataLoader(val_set,batch_size=4,
                                             shuffle=False,num_workers=0)
    #获取图像和标签值
    val_data_iter = iter(val_loader)
    val_image,val_label = val_data_iter.next()

    classes = ('plane','car','bird','cat','deer','dog','frog','horse','ship','truck')
    net = LeNet()
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(),lr=0.001)

    for epoch in range(5):
        running_loss = 0.0#累加损失
        #遍历训练集样本，从零开始
        for step,data in enumerate(train_loader,start=0):
            inputs,labels = data
            optimizer.zero_grad()
            outputs = net(inputs)
            loss = loss_function(outputs,labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if step%500 == 499:
                with torch.no_grad():
                    outputs = net(val_image)
                    predict_y = torch.max(outputs,dim=1)[1]
                    accuracy = torch.eq(predict_y,val_label).sum().item()/val_label.size(0)
                    print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f'%
                          (epoch+1,step+1,running_loss/500,accuracy))
                    running_loss = 0.0

    print('Finished Training')
    save_path = './Lenet.pth'
    torch.save(net.state_dict(),save_path)

if __name__ =='__main__':
    main()
















